import argparse
from typing import List

from pydantic import BaseModel
from tqdm import tqdm
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input-file-path", type=str)
parser.add_argument("-o", "--output-file-path", type=str)
args = parser.parse_args()

# 读取文件内容
with open(args.input_file_path, "r", encoding="utf-8") as file:
    sentences = [line.strip() for line in file if line.strip()]

# sentences = sentences[:1000]

# 按批次处理数据，每批次100条
batch_size = 1000
batches = [sentences[i : i + batch_size] for i in range(0, len(sentences), batch_size)]

# 模型和tokenizer初始化
model_path = "Qwen/Qwen2-7B-Instruct"  # 模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
sampling_params = SamplingParams(temperature=0, top_p=0.9, max_tokens=512)
prefix_cached_llm = LLM(
    model=model_path,
    enable_prefix_caching=True,
    disable_custom_all_reduce=True,
    gpu_memory_utilization=0.9,
    tensor_parallel_size=1,
    max_model_len=4096,
)


class Response(BaseModel):
    keywords: List[str]
    history: List[str] = []

    @classmethod
    def from_str(cls, s: str, history: str):
        start = s.find("{")
        end = s.rfind("}") + 1
        try:
            obj = cls.model_validate_json(s[start:end])
        except Exception:
            obj = Response(keywords=[])
        obj.history = history.split("@")
        if all(["您好" in sentence for sentence in obj.history]):
            obj.keywords = []
        return obj


system_prompt = """
你是智政务，专门为政务话务员设计的智能助手。你能够通过对话记录分析市民的诉求，并准确返回相关的查询关键词。
你的能力有:
- 分析对话内容，提取关键信息。
- 识别市民诉求中的政务知识点。
- 返回对应的查询关键词。
如果没有提取到关键词，或者市民没有回话，你可以返回空列表。

查询关键词的格式为：
{
    "keywords": ["关键词1", "关键词2", "关键词3", ...],
}
"""

example_user_prompt_1 = """呃，我就是我不是报考那个网络预约出租汽车驾驶员证啊，前段时间不是审核，不通过，然后打电话过去问你们吗。@现在还没回复，现在比较着就是比较着急，要报考提交资料了。@这边看到本机号码有相关的一个工单记录，因为承办单位收到工单次日起会在十个工作日内答复的暂时我们也没有收到回应，我们会帮您催促。@对你帮我催一下吧。
"""

example_assistant_prompt_1 = """
这段对话中市民的诉求是关于“网络预约出租汽车驾驶员证”的报考问题，特别是关于审核不通过后的处理进度。
根据市民的诉求内容，以下是相关的查询关键词：
{
    "keywords": ["网络预约出租汽车驾驶员证", "审核不通过原因", "审核结果查询", "工单处理进度", "工单催促", "证书申请资料提交"]
}
"""

example_user_prompt_2 = """您好，请问有什么可以帮您。@您好，电话已经接通，请说。@您好，这边听不到您的一个有效回应。@您好，这边听到您的一个有效回应，如果需要请稍后重拨，再见。"""

example_assistant_prompt_2 = """
这段对话中市民始终没有回复话务员的询问。不清楚市民的诉求是什么。
根据铜壶内容，没有提取到相关的查询关键词。
{
    "keywords": []
}
"""

example_user_prompt_3 = """好的@呃，这边会帮您反馈这个就是说是体外呃，移走之后，呃新加了一块铁板在该地进行这个招生扰民的情情况的。@呃，对对对对，让他们处理一下就是底下垫的东西就可以了，就就不会有有电动车过来就不会很响，要不然我们怎么睡觉，三个半夜有个电车，一个有很大很大的响声。@嗯，好会帮您反馈您的诉求的话呢，有看到我帮您呃，反映要求在贴板下方增设隔音材质的一个物品的。
"""

example_assistant_prompt_3 = """
这段对话中市民的诉求是关于“铁板安装扰民”和“电动车噪音”的问题，特别是要求增设隔音材质。
根据市民的诉求内容，以下是相关的查询关键词：
{
    "keywords": ["铁板安装扰民", "噪音污染投诉", "隔音措施", "城市管理", "噪音治理", "居民环境投诉", "电动车噪音"]
}
"""


# 初始化结果列表
all_results = []


with open(args.output_file_path, "a", encoding="utf-8") as file:
    # 逐批处理数据
    for batch_index, batch in tqdm(enumerate(batches)):
        generating_prompts = []
        for prompt in batch:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": example_user_prompt_1},
                {"role": "assistant", "content": example_assistant_prompt_1},
                {"role": "user", "content": example_user_prompt_2},
                {"role": "assistant", "content": example_assistant_prompt_2},
                {"role": "user", "content": example_user_prompt_3},
                {"role": "assistant", "content": example_assistant_prompt_3},
                {"role": "user", "content": prompt},
            ]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            generating_prompts.append(text)

        # Warmup so that the shared prompt's KV cache is computed.
        if batch_index == 0:
            prefix_cached_llm.generate(generating_prompts[0], sampling_params)

        # Generate with prefix caching.
        outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)

        # Collect results
        batch_results: List[Response] = []
        for history, output in zip(batch, outputs):
            generated_text = output.outputs[0].text
            response = Response.from_str(generated_text, history)
            batch_results.append(response)

        # 将每个批次的结果追加到输出文件中
        for result in batch_results:
            file.write(result.model_dump_json())
            file.write("\n")
            file.flush()
